module aes_key_extend(
    input                   clk,
    input                   rst_n,

    input                   key_vld,
    input   [255:0]         key,
    input   [1:0]           key_mode,
    input   [15:0][255:0]   sbox,

    output  [9:0][255:0]    key_extend,
    output                  extend_done
);

wire    [9:0][255:0]        sub_key;
reg     [9:0][255:0]        key_extend;
reg                         extend_done;

aes_key_extend_sub x_aes_key_extend_sub_0(
    .key(key),
    .sbox(sbox),
    .index(0),

    .sub_key(sub_key[0])
);

generate
genvar i;
for(i=1;i<10;i++)
    aes_key_extend_sub x_aes_key_extend_sub_i(
        .key(key_extend[i-1]),
        .sbox(sbox),
        .index(i),
    
        .sub_key(sub_key[i])
    ); 
endgenerate

always @(posedge clk or negedge rst_n)begin
    if(!rst_n)
        for(integer i=0;i<10;i++)
            key_extend[i] <= 256'b0;
    else if(key_vld)
        for(integer i=0;i<10;i++)
            key_extend[i] <= sub_key[i];
end

always @(posedge clk or negedge rst_n)begin
    if(!rst_n)
        extend_done <= 1'b0;
    else if(key_vld)
        extend_done <= 1'b1;
end



endmodule
